﻿using System;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using Ast.Core.WinApi;

namespace Ast.Core.Services.Extensions
{
    public static class DeviceControllerExtensions
    {
        public static byte[] ReadBytes(this IDeviceController deviceController, IntPtr address, int length)
        {
            var request = new IoCtl.ReadMemoryRequest {Address = address, Length = length};
            byte[] retVal;
            deviceController.Send(IoCtl.IoCtlReadMemory, request, out retVal);
            return retVal;
        }

        public static T ReadAs<T>(this IDeviceController deviceController, IntPtr address)
        {
            var request = new IoCtl.ReadMemoryRequest {Address = address, Length = Marshal.SizeOf<T>()};
            T retVal;
            deviceController.Send(IoCtl.IoCtlReadMemory, request, out retVal);
            return retVal;
        }

        public static T ReadAs<T>(this ReadChain readChain, Func<IntPtr, IntPtr> modifyFunc = null)
        {
            var address = readChain.Value;
            if (modifyFunc != null)
            {
                address = modifyFunc(address);
            }
            return ReadAs<T>(readChain.DeviceController, address);
        }

        public static T[] ReadArrayAs<T>(this IDeviceController deviceController, IntPtr address, int length)
        {
            var size = Marshal.SizeOf<T>();
            var totalSize = size*length;
            var buffer = ReadBytes(deviceController, address, totalSize);
            return buffer.ReadArrayAs<T>(length);

            // old way
//            var size = Marshal.SizeOf<T>();
//            var retVal = new List<T>(length);
//            for (var i = 0; i < length; i++)
//            {
//                var val = ReadAs<T>(deviceController, address + size * i);
//                retVal.Add(val);
//            }
//            return retVal.ToArray();
        }

        public static T[] ReadArrayAs<T>(this IDeviceController deviceController, IntPtr address)
        {
            var size = Marshal.SizeOf<T>();
            var retVal = new List<T>();
            for (var i = 0;; i++)
            {
                var val = ReadAs<T>(deviceController, address + size*i);
                if (val.Equals(default(T)))
                {
                    break;
                }
                retVal.Add(val);
            }
            return retVal.ToArray();
        }

        public static T[] ReadArrayAs<T>(this ReadChain readChain, int length, Func<IntPtr, IntPtr> modifyFunc = null)
        {
            var address = readChain.Value;
            if (modifyFunc != null)
            {
                address = modifyFunc(address);
            }
            return ReadArrayAs<T>(readChain.DeviceController, address, length);
        }

        public static T[] ReadArrayAs<T>(this ReadChain readChain, Func<IntPtr, IntPtr> modifyFunc = null)
        {
            var address = readChain.Value;
            if (modifyFunc != null)
            {
                address = modifyFunc(address);
            }
            return ReadArrayAs<T>(readChain.DeviceController, address);
        }

        public static ReadChain ReadAsIntPtr(this IDeviceController deviceController, IntPtr address)
        {
            var readChain = new ReadChain {DeviceController = deviceController, Value = ReadAs<IntPtr>(deviceController, address)};
            return readChain;
        }

        public static ReadChain ReadAsIntPtr(this ReadChain readChain, Func<IntPtr, IntPtr> modifyFunc = null)
        {
            var address = readChain.Value;
            if (modifyFunc != null)
            {
                address = modifyFunc(address);
            }
            return ReadAsIntPtr(readChain.DeviceController, address);
        }

        public static IntPtr Invoke(this IDeviceController deviceController, string functionName, params IntPtr[] parameters)
        {
            return Invoke<IntPtr>(deviceController, functionName, parameters);
        }

        public static T Invoke<T>(this IDeviceController deviceController, string functionName, params IntPtr[] parameters)
        {
            var request = new IoCtl.InvokeFunctionRequest {FunctionName = functionName, Params = new IntPtr[20], ReturnZise = Marshal.SizeOf<T>()};
            if (parameters != null)
            {
                request.ParamCount = parameters.Length;
                Array.Copy(parameters, request.Params, parameters.Length);
            }
            request.IsReturnValueType = typeof (T) == typeof (IntPtr);

            T retVal;
            deviceController.Send(IoCtl.IoCtlInvokeFunction, request, out retVal);
            return retVal;
        }

        public static string ReadAsStringAnsi(this IDeviceController deviceController, IntPtr address, int length)
        {
            var bytes = ReadBytes(deviceController, address, length);
            return System.Text.Encoding.ASCII.GetString(bytes);
        }

        public static string ReadAsString(this IDeviceController deviceController, IntPtr address, int length)
        {
            var request = new IoCtl.ReadMemoryRequest {Address = address, Length = length};
            string retVal;
            deviceController.Send(IoCtl.IoCtlReadMemory, request, out retVal);
            return retVal;
        }

        public static string ReadAsUnicodeString(this IDeviceController deviceController, Structure.UnicodeString unicodeString)
        {
            return unicodeString.Length <= 0
                ? null
                : ReadAsString(deviceController, unicodeString.Buffer, unicodeString.Length);
        }

        public static string ReadAsUnicodeString(this IDeviceController deviceController, IntPtr address)
        {
            var unicodeString = ReadAs<Structure.UnicodeString>(deviceController, address);
            return ReadAsUnicodeString(deviceController, unicodeString);
        }

        public static KernelMemoryBlock AllocateMemory(this IDeviceController deviceController, int size)
        {
            var address = Invoke(deviceController, "ExAllocatePoolWithTag", IntPtr.Zero, new IntPtr(size), new IntPtr(IoCtl.MemTag));
            Invoke(deviceController, "RtlZeroMemory", address, new IntPtr(size));
            return new KernelMemoryBlock {DeviceController = deviceController, Value = address};
        }

        public static void FreeMemory(this IDeviceController deviceController, IntPtr address)
        {
            Invoke(deviceController, "ExFreePoolWithTag", address, new IntPtr(IoCtl.MemTag));
        }

        public static void WriteBytes(this IDeviceController deviceController, IntPtr address, byte[] data)
        {
            var request = new IoCtl.WriteMemoryRequest {Address = address, Length = data.Length};
            var inDataSize = Marshal.SizeOf<IoCtl.WriteMemoryRequest>() + data.Length;
            var inDataBuffer = Marshal.AllocHGlobal(inDataSize);
            Marshal.StructureToPtr(request, inDataBuffer, true);
            Marshal.Copy(data, 0, inDataBuffer + Marshal.SizeOf<IoCtl.WriteMemoryRequest>(), data.Length);
            var returnDataSize = default(uint);
            deviceController.Send(IoCtl.IoCtlWriteMemory, inDataBuffer, (uint) inDataSize, IntPtr.Zero, 0, ref returnDataSize);
            Marshal.FreeHGlobal(inDataBuffer);
        }

        public static void WriteAs<T>(this IDeviceController deviceController, IntPtr address, T value) where T : struct
        {
            var size = Marshal.SizeOf<T>();
            var data = new byte[size];
            var buffer = Marshal.AllocHGlobal(size);
            Marshal.StructureToPtr(value, buffer, true);
            Marshal.Copy(buffer, data, 0, size);
            Marshal.FreeHGlobal(buffer);
            WriteBytes(deviceController, address, data);

            // old way
//            var request = new IoCtl.WriteMemoryRequest {Address = address, Length = Marshal.SizeOf<T>()};
//            var inDataSize = Marshal.SizeOf<IoCtl.WriteMemoryRequest>() + Marshal.SizeOf<T>();
//            var inDataBuffer = Marshal.AllocHGlobal(inDataSize);
//            Marshal.StructureToPtr(request, inDataBuffer, true);
//            Marshal.StructureToPtr(value, inDataBuffer + Marshal.SizeOf<IoCtl.WriteMemoryRequest>(), true);
//            var returnDataSize = default(uint);
//            deviceController.Send(IoCtl.IoCtlWriteMemory, inDataBuffer, (uint) inDataSize, IntPtr.Zero, 0, ref returnDataSize);
//            Marshal.FreeHGlobal(inDataBuffer);
        }

        public static KernelMemoryBlock AllocateMemoryForUnicodeString(this IDeviceController deviceController, int maximumLength)
        {
            var mem = AllocateMemory(deviceController, Marshal.SizeOf<Structure.UnicodeString>() + maximumLength);
            deviceController.WriteAs(mem.Value, new Structure.UnicodeString {Length = 0, MaximumLength = (ushort) maximumLength, Buffer = mem.Value + Marshal.SizeOf<Structure.UnicodeString>()});
            return mem;
        }
    }

    public class ReadChain
    {
        public IDeviceController DeviceController { get; set; }
        public IntPtr Value { get; set; }
    }
}